"""
合同要素抽取
"""
import json
from typing import Any, Dict
from core.llm_answer import answer_with_llm
from core.milvus_vector import MilvusVector
from core.document import Document
from core.config import settings
from core.reranker import ReRanker
from loguru import logger
import rapidfuzz
import re


def extract_prompt(contract_text, rule_name, extract_hint) -> str:

    if extract_hint != '':
        extract_text = f'按照如下指示：`{extract_hint}`，并根据合同内容排版的特点，准确找出`{rule_name}`所意指的文本内容'
    else:
        extract_text = f'根据合同内容排版的特点，准确找出`{rule_name}`所意指的文本内容'

    if rule_name.find('合同名称') != -1:
        extract_text += f'，协议也是合同的一种，通常，`{rule_name}`位于合同的开始部分'

    if re.search('人[^\u4E00-\u9F5A]*$', rule_name) is not None:
        extract_text += f'，注意：所找出的文本内容必须包含有效的人名'

    prompt = f'''
      忘掉之前的所有设定和对话，从现在开始，你是一名专业的合同审查助手。现合同中有以下条款段落(按顺序排列)：
      `{contract_text}`

      你需要基于你的专业知识，{extract_text}，
      结果以 JSON 格式输出，包含`content` 字段，字段说明如下：
      - content：`{rule_name}`所指向的文本内容，如果未找到或所指向的内容未填写，则为空字符串。
    '''
    return prompt


async def rule_extract_api(fileId: str, rule: Dict[str, Any]) -> Dict:
    """
    合同要素的内容抽取
    """
    info = {'fileId': fileId, 'rule': rule}
    logger.info(f'要素内容抽取: {json.dumps(info, ensure_ascii=False)}')

    # 召回要素相关的片段
    milvus = MilvusVector()
    rule_name = rule['name']
    extract_hint = rule.get('extractHint', '') or ''

    if extract_hint != '':
        query = f'按照如下指示：`{extract_hint}`，找出`{rule_name}`相关的文本内容。'
    else:
        query = rule_name
    documents = milvus.hybrid_search(fileId, query, top_k=30)
    reranker = ReRanker()
    documents = reranker.rerank(query, documents, settings.TOP_K)   # 文档重排

    content, related_dict = await answer_by_documents(documents, rule_name, extract_hint)
    if content == '':
        # 如果要素抽取失败，则选取合同的前N个片段重新抽取
        all_documents = milvus.fetch_documents(fileId)
        documents = all_documents[0:30]
        content, related_dict = await answer_by_documents(documents, rule_name, extract_hint)
        if content == '':
            # 如果要素抽取还是失败，则选取合同的后N个片段重新抽取
            documents = all_documents[-1:-30:-1]
            content, related_dict = await answer_by_documents(documents, rule_name, extract_hint)

    if content != '' and len(documents) != 0:
        partial_ratios = dict()
        key2id = dict()
        for key, block in related_dict.items():
            partial_ratios[key] = rapidfuzz.fuzz.partial_ratio(block['content'], content)
            key2id[key] = block['id']
        sored_ratios = sorted(partial_ratios.items(), key=lambda x: -x[1])
        paragraph_id = key2id[sored_ratios[0][0]]
        sorted_ids = list(dict.fromkeys([item[0] for item in sored_ratios]))
    else:
        paragraph_id = None
        sorted_ids = list(dict.fromkeys([doc.metadata['id'] for doc in documents]))

    # 返回结果
    extract_result = {
        'fileId': fileId,             # 合同文件id
        'id': rule.get('id', None),   # 要素id
        'name': rule_name,            # 要素名称
        'extractHint': extract_hint,  # 要素抽取的hint
        'related_ids': sorted_ids,    # 召回的相关段落id
        'content': content,           # 要素的内容
        'paragraph_id': paragraph_id, # 要素的内容关联的段落id
    }
    logger.info(f'要素抽取结果: {json.dumps(extract_result, ensure_ascii=False)}')
    return extract_result


async def answer_by_documents(documents: list[Document], rule_name:str, extract_hint: str):
    """从给定的Documents中抽取要素"""
    if settings.DEBUG:
        logger.debug('召回的文本片段：')
        for doc in documents:
            logger.debug(doc)

    # 调用大模型进行要素抽取
    prompt = None
    related_dict = None
    for k in reversed(range(0, len(documents))):
        related_texts, related_dict = reorganize(documents[0: k])
        prompt = extract_prompt('\n'.join(related_texts), rule_name, extract_hint)
        if len(prompt) < settings.LLM_MAX_INPUT:
            break

    answer = await answer_with_llm(prompt, parse_type='dict')
    if settings.DEBUG:
        logger.debug(f'prompt：{prompt}')
        logger.debug(f'answer：{answer}')

    # 与要素最相关的内容
    content = answer.get('content', '')
    content = re.sub('###?\s+', '', content)
    return content, related_dict

def reorganize(documents) :
    # 按合同原始顺序拼接成markdown形式的文本, 标题也需整合进去
    retrieved = dict()
    related_dict = dict()
    for doc in documents:
        related_blocks = doc.metadata['related_blocks']
        for key, block in related_blocks.items():
            start_index = int(key.split('-')[1])
            content = block['content']

            block_index = int(key.split('-')[0])
            if block_index not in retrieved:
                retrieved[block_index] = dict()
            retrieved[block_index][start_index] = content
            related_dict[key] = block

    retrieved = sorted(retrieved.items(), key=lambda x: x[0])
    related_texts = []
    for block_index, content_items in retrieved:
        content_items = sorted(content_items.items(), key=lambda x: x[0])
        related_texts.append(''.join([content for _, content in content_items]))

    return related_texts, related_dict